import base64
from collections import defaultdict
import difflib
from io import BytesIO
import os
import shutil
from tkinter import Image

import numpy as np
import requests
from src.modules.image_processor import ImageProcessor
from src.modules.ocr_processor import OCRProcessor
from src.utils.logger import log_error, log_info


class WatermarkDetector:
    def __init__(self):
        """
        初始化水印检测器
        :param ocr_processor: OCR处理器实例
        :param image_processor: 图像处理器实例
        """
        self._ocr_processor = None
        self._image_processor = None

    @property
    def image_processor(self):
        if self._image_processor is None:
            from src.modules.image_processor import ImageProcessor
            self._image_processor = ImageProcessor()
        return self._image_processor

    @property
    def ocr_processor(self):
        if self._ocr_processor is None:
            from src.modules.ocr_processor import OCRProcessor
            self._ocr_processor = OCRProcessor()
        return self._ocr_processor

    def detect_watermark_with_ocr(self, folder_path, min_repeat=8):
        """
        分析文件夹中图像OCR结果，判断水印和字幕的类型
        :param folder_path: 图像文件夹路径
        :param min_repeat: 文本重复出现的最小次数
        :return: 检测结果字典
        """
        if not os.path.isdir(folder_path):
            log_info(f"指定的路径不是有效的文件夹: {folder_path}")
            return {"has_watermark": False, "type": "none", "texts": []}

        log_info(f"开始检测文件夹中的图片水印: {folder_path}")

        # 存储文本出现情况
        text_occurrences = defaultdict(list)  # {text: [position_info1, position_info2, ...]}
        image_info = {}  # 存储图像基本信息
        image_count = 0
        text_count = 0
        
        # 图像区域文本追踪 - 用于不同类型的水印/字幕检测
        top_region_texts = defaultdict(list)     # 顶部区域文本 (y < 0.3)
        center_region_texts = defaultdict(list)  # 中部区域文本 (0.3 <= y <= 0.7)
        bottom_region_texts = defaultdict(list)  # 底部区域文本 (y > 0.7)
        
        # 跟踪每张图片的文本总高度
        text_heights_by_image = defaultdict(float)
        
        # 1. 收集所有图像的OCR结果和位置信息
        for file_name in os.listdir(folder_path):
            file_path = os.path.join(folder_path, file_name)
            if os.path.isfile(file_path) and file_name.lower().endswith(('.jpg')):
                try:
                    # 读取图像获取尺寸
                    img = self.image_processor.imread_unicode(file_path)
                    if img is None:
                        continue

                    height, width, _ = img.shape
                    image_info[file_name] = {"width": width, "height": height, "area": width * height}
                    image_count += 1

                    # OCR识别
                    detected = self.ocr_processor.recognize_filtered_subtitles(file_path)
                    
                    # 此图片中的文本总高度
                    total_height_this_image = 0
                    
                    for item in detected:
                        text = item["text"].strip()
                        if not text or (text.isdigit() and len(text) <= 3):
                            continue
                        text_count += 1
                        position = item["position"]

                        # 找到相似文本的键
                        matched_key = self.find_similar_key(text_occurrences, text)
                        occurrence_key = matched_key if matched_key else text
                        
                        # 计算相对位置和尺寸
                        rel_x = position["center_x"] / width
                        rel_y = position["center_y"] / height
                        rel_width = position["width"] / width
                        rel_height = position["height"] / height
                        
                        # 累计这个图片中的文本高度
                        total_height_this_image += rel_height
                        
                        # 存储位置信息
                        occurrence_info = {
                            "file": file_name,
                            "text": text,
                            "position": position,
                            "rel_x": rel_x,
                            "rel_y": rel_y,
                            "rel_width": rel_width,
                            "rel_height": rel_height
                        }
                        
                        # 通用文本出现记录
                        text_occurrences[occurrence_key].append(occurrence_info)
                        
                        # 按区域记录文本出现
                        if rel_y < 0.3:
                            top_region_texts[file_name].append(occurrence_info)
                        elif rel_y > 0.7:
                            bottom_region_texts[file_name].append(occurrence_info)
                        else:
                            center_region_texts[file_name].append(occurrence_info)

                    # 记录这个图片的文本总高度
                    text_heights_by_image[file_name] = total_height_this_image

                except Exception as e:
                    log_info(f"识别图片 {file_name} 时出错: {e}")
                    log_error(f"识别图片 {file_name} 时出错: {e}")

        # 2. 分析是否存在水印或字幕
        log_info(f"分析识别结果，共处理 {image_count} 张图像，识别到 {text_count} 个文本项")

        # 尝试删除临时文件夹
        try:
            shutil.rmtree(folder_path)
        except Exception as e:
            log_info(f"删除临时文件夹时出错: {e}")

        # 如果没有处理任何图像，返回无水印
        if image_count == 0:
            return {"has_watermark": False, "type": "none", "texts": []}

        # 初始化检测结果
        detection_results = {
            "固定水印": {"detected": False, "confidence": 0, "texts": []},
            "移动水印": {"detected": False, "confidence": 0, "texts": []},
            "居中字幕": {"detected": False, "confidence": 0, "count": 0},
            "大块字幕": {"detected": False, "confidence": 0, "count": 0},
            "普通字幕": {"detected": False, "confidence": 0, "count": 0}
        }

        # 3. 检测五种不同类型的水印/字幕
        
        # 3.1 检测固定水印和移动水印 (基于特定文本的出现方式)
        for text, occurrences in text_occurrences.items():
            # 只分析出现足够次数的文本
            unique_images = len(set(occ["file"] for occ in occurrences))
            if unique_images < min_repeat or unique_images < image_count * 0.4:
                continue
                
            # 计算位置和尺寸稳定性
            rel_positions = [(occ["rel_x"], occ["rel_y"]) for occ in occurrences]
            var_x = np.var([pos[0] for pos in rel_positions])
            var_y = np.var([pos[1] for pos in rel_positions])
            position_variance = var_x + var_y
            size_variance = np.var([occ["rel_width"] * occ["rel_height"] for occ in occurrences])
            
            # 判断是否是固定水印 (位置固定)
            is_stable_position = position_variance < 0.05  # 位置变化不大
            is_stable_size = size_variance < 0.01  # 大小变化不大
            
            # 计算文本出现在不同区域的次数
            top_occurrences = sum(1 for occ in occurrences if occ["rel_y"] < 0.3)
            center_occurrences = sum(1 for occ in occurrences if 0.3 <= occ["rel_y"] <= 0.7)
            bottom_occurrences = sum(1 for occ in occurrences if occ["rel_y"] > 0.7)
            
            # 文本长度限制 (水印通常较短)
            if len(text) > 20:
                continue
                
            # 判断固定水印
            if is_stable_position and is_stable_size:
                confidence = 0.8 + (0.2 * (1 - position_variance))  # 0.8-1.0的置信度
                log_info(f"检测到固定水印: '{text}', 位置方差: {position_variance:.4f}, 置信度: {confidence:.2f}")
                detection_results["固定水印"]["detected"] = True
                detection_results["固定水印"]["texts"].append(text)
                if confidence > detection_results["固定水印"]["confidence"]:
                    detection_results["固定水印"]["confidence"] = confidence
            
            # 判断移动水印 (同一文本出现在多个不同位置)
            elif unique_images >= min_repeat and position_variance >= 0.05:
                # 如果文本主要出现在底部，可能是普通字幕而非移动水印
                if bottom_occurrences > 0.7 * len(occurrences):
                    continue
                    
                confidence = 0.7 + (0.2 * (position_variance - 0.05) / 0.45)  # 0.7-0.9的置信度
                log_info(f"检测到移动水印: '{text}', 位置方差: {position_variance:.4f}, 置信度: {confidence:.2f}")
                detection_results["移动水印"]["detected"] = True
                detection_results["移动水印"]["texts"].append(text)
                if confidence > detection_results["移动水印"]["confidence"]:
                    detection_results["移动水印"]["confidence"] = confidence
        
        # 3.2 检测居中字幕 (不需要是相同文本)
        center_text_images = [file for file, texts in center_region_texts.items() if texts]
        center_text_ratio = len(center_text_images) / image_count if image_count > 0 else 0

        # 增加min_repeat限制
        if len(center_text_images) >= min_repeat and center_text_ratio >= 0.4:
            confidence = 0.7 + (0.2 * center_text_ratio)
            total_center_texts = sum(len(texts) for texts in center_region_texts.values())
            log_info(f"检测到居中字幕: {total_center_texts} 处, {center_text_ratio:.2f} 的图片有中心文本, 置信度: {confidence:.2f}")
            detection_results["居中字幕"]["detected"] = True
            detection_results["居中字幕"]["confidence"] = confidence
            detection_results["居中字幕"]["count"] = total_center_texts

        # 3.3 检测大块字幕 (文本高度占比大)
        large_text_images = [file for file, height in text_heights_by_image.items() if height >= 0.2]
        large_text_ratio = len(large_text_images) / image_count if image_count > 0 else 0

        # 增加min_repeat限制
        if len(large_text_images) >= min_repeat and large_text_ratio >= 0.3:
            confidence = 0.7 + (0.3 * large_text_ratio)
            log_info(f"检测到大块字幕: {len(large_text_images)} 张图片有大块文本, 覆盖比例: {large_text_ratio:.2f}, 置信度: {confidence:.2f}")
            detection_results["大块字幕"]["detected"] = True
            detection_results["大块字幕"]["confidence"] = confidence
            detection_results["大块字幕"]["count"] = len(large_text_images)

        # 3.4 检测普通字幕 (顶部或底部文本)
        edge_text_images = set()
        for file in top_region_texts:
            if top_region_texts[file]:
                edge_text_images.add(file)
        for file in bottom_region_texts:
            if bottom_region_texts[file]:
                edge_text_images.add(file)
                
        edge_text_ratio = len(edge_text_images) / image_count if image_count > 0 else 0

        # 增加min_repeat限制
        if len(edge_text_images) >= min_repeat and edge_text_ratio >= 0.5 and not detection_results["移动水印"]["detected"] and not detection_results["固定水印"]["detected"]:
            confidence = 0.6 + (0.3 * edge_text_ratio)
            total_edge_texts = sum(len(texts) for texts in top_region_texts.values()) + sum(len(texts) for texts in bottom_region_texts.values())
            log_info(f"检测到普通字幕: {total_edge_texts} 处, {edge_text_ratio:.2f} 的图片有边缘文本, 置信度: {confidence:.2f}")
            detection_results["普通字幕"]["detected"] = True
            detection_results["普通字幕"]["confidence"] = confidence
            detection_results["普通字幕"]["count"] = total_edge_texts
        
        # 4. 确定最终结果 (按优先级排序)
        # 设置优先级
        priorities = {
            "移动水印": 2,
            "大块字幕": 3,
            "居中字幕": 2,
            "固定水印": 1,
            "普通字幕": 1
        }

        # 初始化最终结果
        final_type = None
        final_texts = []
        highest_priority = 0
        highest_confidence = 0

        # 根据优先级和置信度确定最终结果
        for type_name, info in detection_results.items():
            if info["detected"]:
                current_priority = priorities.get(type_name, 0)
                current_confidence = info.get("confidence", 0)
                
                log_info(f"考虑类型: {type_name}, 优先级: {current_priority}, 置信度: {current_confidence}")
                
                # 如果优先级更高或优先级相同但置信度更高，则更新结果
                if (current_priority > highest_priority) or (current_priority == highest_priority and current_confidence > highest_confidence):
                    highest_priority = current_priority
                    highest_confidence = current_confidence
                    
                    # 确定返回类型 (大块字幕和居中字幕都归为"移动水印"类型)
                    if type_name in ["大块字幕", "居中字幕"]:
                        final_type = "移动水印"
                        final_texts = [type_name]  # 使用类型名称作为文本描述
                    elif type_name == "普通字幕":
                        final_type = "字幕"
                        final_texts = ["普通字幕"]
                    else:
                        final_type = type_name
                        final_texts = info.get("texts", []) if "texts" in info else [type_name]

        # 如果没有任何检测结果，但文本数量足够，则认为是普通字幕
        if not final_type and text_count >= min_repeat:
            final_type = "字幕"
            final_texts = ["检测到字幕"]
        # 如果仍然没有结果，则认为没有水印
        elif not final_type:
            return {"has_watermark": False, "type": "none", "texts": []}

        log_info(f"最终检测结果: {final_type}, 文本: {', '.join(final_texts if final_texts else ['未提取文本'])}, 优先级: {highest_priority}, 置信度: {highest_confidence:.2f}")


        # 打印检测总结信息
        log_info("==== 水印检测总结 ====")
        log_info(f"总计分析图片数: {image_count}，识别文本数: {text_count}")
        
        # 固定水印总结
        if detection_results["固定水印"]["detected"]:
            texts = detection_results["固定水印"]["texts"]
            log_info(f"固定水印文本: {', '.join(texts) if texts else '无文本内容'}")
            log_info(f"固定水印置信度: {detection_results['固定水印']['confidence']:.2f}")
        
        # 移动水印总结
        if detection_results["移动水印"]["detected"]:
            texts = detection_results["移动水印"]["texts"]
            log_info(f"移动水印文本: {', '.join(texts) if texts else '无文本内容'}")
            log_info(f"移动水印置信度: {detection_results['移动水印']['confidence']:.2f}")
        
        # 居中字幕总结
        if detection_results["居中字幕"]["detected"]:
            count = detection_results["居中字幕"]["count"]
            images_count = len(center_text_images)
            log_info(f"居中字幕: 共{count}处，出现在{images_count}张图片中（{center_text_ratio:.1%}的图片）")
            log_info(f"居中字幕置信度: {detection_results['居中字幕']['confidence']:.2f}")
            
            # 打印部分样例内容
            sample_texts = []
            for file, texts in center_region_texts.items():
                if texts:
                    sample_texts.extend([text["text"] for text in texts[:2]])  # 每个图片最多取2个文本
                if len(sample_texts) >= 5:  # 最多展示5个样例
                    break
            if sample_texts:
                log_info(f"居中字幕样例: {', '.join(sample_texts[:5])}")
        
        # 大块字幕总结
        if detection_results["大块字幕"]["detected"]:
            count = detection_results["大块字幕"]["count"]
            log_info(f"大块字幕: 出现在{count}张图片中（{large_text_ratio:.1%}的图片）")
            log_info(f"大块字幕置信度: {detection_results['大块字幕']['confidence']:.2f}")
        
        # 普通字幕总结
        if detection_results["普通字幕"]["detected"]:
            total_edge_texts = sum(len(texts) for texts in top_region_texts.values()) + sum(len(texts) for texts in bottom_region_texts.values())
            top_count = sum(len(texts) for texts in top_region_texts.values())
            bottom_count = sum(len(texts) for texts in bottom_region_texts.values())
            log_info(f"普通字幕: 共{total_edge_texts}处，顶部{top_count}处，底部{bottom_count}处")
            log_info(f"普通字幕出现在{len(edge_text_images)}张图片中（{edge_text_ratio:.1%}的图片）")
            log_info(f"普通字幕置信度: {detection_results['普通字幕']['confidence']:.2f}")
            
            # 打印部分样例内容
            sample_texts = []
            for collection in [top_region_texts, bottom_region_texts]:
                for file, texts in collection.items():
                    if texts:
                        sample_texts.extend([text["text"] for text in texts[:1]])  # 每个图片最多取1个文本
                    if len(sample_texts) >= 5:  # 最多展示5个样例
                        break
            if sample_texts:
                log_info(f"边缘字幕样例: {', '.join(sample_texts[:5])}")
        
        # 最终结果总结
        log_info(f"最终判定类型: {final_type}")
        log_info(f"检测到的文本: {', '.join(final_texts)}")
        log_info(f"最终置信度: {highest_confidence:.2f}")
        log_info("==== 水印检测结束 ====")

        return {
            "has_watermark": True,
            "type": final_type,
            "texts": final_texts,
            "priority": highest_priority,
            "confidence": highest_confidence,
            "detailed_results": detection_results  # 包含详细的检测结果
        }

    def find_similar_key(self, text_occurrences, text, threshold=0.8):
        for existing_text in text_occurrences.keys():
            ratio = difflib.SequenceMatcher(None, existing_text, text).ratio()
            if ratio > threshold:
                return existing_text
        return None    

    def detect_watermark_with_huggingface(self, image_path, api_key,
                                          model_name="prithivMLmods/Watermark-Detection-SigLIP2"):
        """
        使用 Hugging Face 的模型对文本进行摘要
        :param text: 要摘要的文本
        :param model_name: 使用的模型名称（默认为 "prithivMLmods/Watermark-Detection-SigLIP2"）
        :return: 摘要后的文本
        """
        with Image.open(image_path) as img:
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        url = f"https://api-inference.huggingface.co/models/{model_name}"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "inputs": image_base64
        }
        response = requests.post(url, headers=headers, json=payload)
        if not response.ok:
            raise RuntimeError(f"调用 Hugging Face 模型时出错: {response.text}")

        result = response.json()
        watermark_score = next((item['score'] for item in result if item['label'] == 'Watermark'), 0)
        no_watermark_score = next((item['score'] for item in result if item['label'] == 'No Watermark'), 0)

        return watermark_score > no_watermark_score, watermark_score, no_watermark_score

